import numpy as np
import math
import pygad
import pandas as pd
import matplotlib.pyplot as plt
import random

plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
plt.rcParams['axes.unicode_minus'] = False

# -------------------------- 1. 基础参数与复用函数 --------------------------
target_true = np.array([0, 200, 0])  # 真目标
r, h = 7, 10  # 真目标的半径7m、高10m
target_fake = np.array([0, 0, 0])  # 假目标
missile_M1 = np.array([20000, 0, 2000])  # M1初始位置
missile_speed = 300  # 导弹速度
drone_FY1 = np.array([17800, 0, 1800])  # FY1初始位置
g = 9.8  # 重力加速度
smoke_radius = 10  # 烟幕半径
smoke_sink_speed = 3  # 下沉速度
smoke_valid_time = 20  # 有效时间
drone_speed_range = [70, 140]  # 无人机速度约束

# 计算导弹飞到假目标的时间（关键约束）
missile_distance_to_fake = np.linalg.norm(target_fake - missile_M1)
missile_time_to_fake = missile_distance_to_fake / missile_speed  # 导弹到达假目标的时间
print(f"导弹到达假目标时间: {missile_time_to_fake:.2f}s")


# 定义获取长方体顶点的函数
def get_cylinder_bounding_box_vertices(target_center, radius, height):
    x_min = target_center[0] - radius
    x_max = target_center[0] + radius
    y_min = target_center[1] - radius
    y_max = target_center[1] + radius
    z_min = target_center[2]
    z_max = target_center[2] + height
    # 长方体顶点
    vertices = [
        np.array([x_min, y_min, z_min]),
        np.array([x_max, y_min, z_min]),
        np.array([x_max, y_max, z_min]),
        np.array([x_min, y_max, z_min]),
        np.array([x_min, y_min, z_max]),
        np.array([x_max, y_min, z_max]),
        np.array([x_max, y_max, z_max]),
        np.array([x_min, y_max, z_max])
    ]
    return vertices


# 预计算目标顶点
target_vertices = get_cylinder_bounding_box_vertices(target_true, r, h)


def missile_real_time_position(missile_init_pos, t):
    # 导弹到达假目标后不再移动
    if t >= missile_time_to_fake:
        return target_fake
    if t < 0:
        return missile_init_pos
    dir_vec = target_fake - missile_init_pos
    dir_norm = np.linalg.norm(dir_vec)
    if dir_norm <= 1e-6:
        return missile_init_pos
    dir_unit_vec = dir_vec / dir_norm
    return missile_init_pos + missile_speed * t * dir_unit_vec


# -------------------------- 2. 烟幕实时位置计算 --------------------------
def get_smoke_explosion_center(drone_init_pos, drone_speed, theta, t_drop, t_detonate):
    # 投放点坐标
    drop_x = drone_init_pos[0] + drone_speed * math.cos(theta) * t_drop
    drop_y = drone_init_pos[1] + drone_speed * math.sin(theta) * t_drop
    drop_z = drone_init_pos[2]
    # 烟幕弹平抛运动
    explosion_x = drop_x + drone_speed * math.cos(theta) * t_detonate
    explosion_y = drop_y + drone_speed * math.sin(theta) * t_detonate
    # 竖直自由下落
    fall_distance = 0.5 * g * (t_detonate ** 2)
    explosion_z = max(drop_z - fall_distance, 0.0)
    return np.array([explosion_x, explosion_y, explosion_z])


def smoke_real_time_center(drone_init_pos, drone_speed, theta, t_drop, t_detonate, t):
    detonate_t = t_drop + t_detonate
    if t < detonate_t - 1e-6:
        return np.array([0, 0, 0])
    explosion_center = get_smoke_explosion_center(drone_init_pos, drone_speed, theta, t_drop, t_detonate)
    sink_distance = smoke_sink_speed * (t - detonate_t)
    real_z = max(explosion_center[2] - sink_distance, 0.0)
    return np.array([explosion_center[0], explosion_center[1], real_z])


# -------------------------- 3. 遮蔽判断 --------------------------
def is_segment_sphere_intersect(missile_pos, vertex_pos, smoke_center, smoke_radius):
    vec_AC = smoke_center - missile_pos
    vec_AB = vertex_pos - missile_pos
    ab_sq_mag = np.dot(vec_AB, vec_AB)
    if ab_sq_mag < 1e-12:
        return False
    k = np.dot(vec_AC, vec_AB) / ab_sq_mag
    if k <= 0.0:
        min_dist = np.linalg.norm(vec_AC)
    elif 0.0 < k < 1.0:
        cross_prod = np.cross(vec_AC, vec_AB)
        min_dist = np.linalg.norm(cross_prod) / np.sqrt(ab_sq_mag)
    else:
        vec_BC = smoke_center - vertex_pos
        min_dist = np.linalg.norm(vec_BC)
    return min_dist <= smoke_radius + 1e-6


def is_effective_occlusion(t, drone_speed, theta, t_drop, t_detonate):
    detonate_t = t_drop + t_detonate
    smoke_start = detonate_t
    smoke_end = detonate_t + smoke_valid_time
    if t < smoke_start - 1e-6 or t > smoke_end + 1e-6:
        return False

    missile_pos = missile_real_time_position(missile_M1, t)
    smoke_center = smoke_real_time_center(drone_FY1, drone_speed, theta, t_drop, t_detonate, t)

    # 简化的遮蔽判断：只检查导弹与烟幕中心的距离
    distance = np.linalg.norm(missile_pos - smoke_center)
    return distance <= smoke_radius + 5  # 增加5米容差


# -------------------------- 4. 多弹总遮蔽时长计算 --------------------------
def calculate_total_duration(drone_speed, theta, t_drop1, t_fuse1, t_drop2, t_fuse2, t_drop3, t_fuse3, time_step=0.1):
    # 为三枚弹分别计算有效时间区间
    intervals = []

    # 第一枚弹
    detonate_t1 = t_drop1 + t_fuse1
    smoke_start1 = detonate_t1
    smoke_end1 = detonate_t1 + smoke_valid_time
    t = smoke_start1
    while t <= smoke_end1 + 1e-6:
        if is_effective_occlusion(t, drone_speed, theta, t_drop1, t_fuse1):
            intervals.append((t, t + time_step))
        t += time_step

    # 第二枚弹
    detonate_t2 = t_drop2 + t_fuse2
    smoke_start2 = detonate_t2
    smoke_end2 = detonate_t2 + smoke_valid_time
    t = smoke_start2
    while t <= smoke_end2 + 1e-6:
        if is_effective_occlusion(t, drone_speed, theta, t_drop2, t_fuse2):
            intervals.append((t, t + time_step))
        t += time_step

    # 第三枚弹
    detonate_t3 = t_drop3 + t_fuse3
    smoke_start3 = detonate_t3
    smoke_end3 = detonate_t3 + smoke_valid_time
    t = smoke_start3
    while t <= smoke_end3 + 1e-6:
        if is_effective_occlusion(t, drone_speed, theta, t_drop3, t_fuse3):
            intervals.append((t, t + time_step))
        t += time_step

    # 合并重叠区间
    if not intervals:
        return 0.0

    intervals.sort(key=lambda x: x[0])
    merged = [intervals[0]]

    for current in intervals[1:]:
        last = merged[-1]
        if current[0] <= last[1]:
            merged[-1] = (last[0], max(last[1], current[1]))
        else:
            merged.append(current)

    # 计算总时长
    total_duration = sum(end - start for start, end in merged)
    return round(total_duration, 3)


# -------------------------- 5. 适应度函数（核心优化逻辑） --------------------------
def fitness_func(ga_instance, solution, solution_idx):
    drone_speed, theta, t_drop1, t_fuse1, t_drop2, t_fuse2, t_drop3, t_fuse3 = solution

    # 温和的罚函数权重
    penalty = 0

    # 速度约束
    if not (drone_speed_range[0] <= drone_speed <= drone_speed_range[1]):
        penalty += 10 * abs(drone_speed - np.clip(drone_speed, drone_speed_range[0], drone_speed_range[1]))

    # 投放间隔约束
    if t_drop2 < t_drop1 + 1:
        penalty += 10 * (t_drop1 + 1 - t_drop2)
    if t_drop3 < t_drop2 + 1:
        penalty += 10 * (t_drop2 + 1 - t_drop3)

    # 起爆时间约束（必须在导弹到达假目标之前）
    detonate_t1 = t_drop1 + t_fuse1
    detonate_t2 = t_drop2 + t_fuse2
    detonate_t3 = t_drop3 + t_fuse3

    if detonate_t1 > missile_time_to_fake:
        penalty += 10 * (detonate_t1 - missile_time_to_fake)
    if detonate_t2 > missile_time_to_fake:
        penalty += 10 * (detonate_t2 - missile_time_to_fake)
    if detonate_t3 > missile_time_to_fake:
        penalty += 10 * (detonate_t3 - missile_time_to_fake)

    # 计算总遮蔽时长
    total_duration = calculate_total_duration(drone_speed, theta, t_drop1, t_fuse1, t_drop2, t_fuse2, t_drop3, t_fuse3)

    # 返回适应度（总遮蔽时长减去罚函数）
    fitness = total_duration - penalty

    # 打印调试信息（每50个个体打印一次）
    if solution_idx % 50 == 0:
        print(f"个体 {solution_idx}: 速度={drone_speed:.2f}, 角度={theta:.2f}, "
              f"投放时间=({t_drop1:.2f}, {t_drop2:.2f}, {t_drop3:.2f}), "
              f"起爆时间=({t_fuse1:.2f}, {t_fuse2:.2f}, {t_fuse3:.2f}), "
              f"遮蔽时长={total_duration:.2f}, 罚分={penalty:.2f}, 适应度={fitness:.2f}")

    return max(fitness, 0)  # 确保适应度不为负


# -------------------------- 6. 初始种群生成函数 --------------------------
def create_starting_population():
    population = []
    for _ in range(100):  # 生成100个初始个体
        # 随机生成参数
        drone_speed = random.uniform(70, 140)
        theta = random.uniform(0, 2 * math.pi)

        # 生成满足投放间隔约束的时间
        max_drop_time = missile_time_to_fake * 0.8  # 预留20%的时间用于起爆
        t_drop1 = random.uniform(0, max_drop_time / 3)
        t_drop2 = random.uniform(t_drop1 + 1, max_drop_time / 2)
        t_drop3 = random.uniform(t_drop2 + 1, max_drop_time)

        # 生成起爆时间，确保在导弹到达之前
        max_fuse_time = missile_time_to_fake - max(t_drop1, t_drop2, t_drop3)
        t_fuse1 = random.uniform(0, max_fuse_time * 0.5)
        t_fuse2 = random.uniform(0, max_fuse_time * 0.7)
        t_fuse3 = random.uniform(0, max_fuse_time)

        population.append([drone_speed, theta, t_drop1, t_fuse1, t_drop2, t_fuse2, t_drop3, t_fuse3])

    return np.array(population)


# -------------------------- 7. 用遗传算法训练 --------------------------
# 用于保存每代最大遮蔽时长
generation_best_fitness = []


def on_generation(ga_instance):
    # 获取当前种群最大适应度
    best_fitness = np.max(ga_instance.last_generation_fitness)
    generation_best_fitness.append(best_fitness)
    print(f"当前迭代次数：{ga_instance.generations_completed}, 最大适应度：{best_fitness:.2f}")


# -------------------------- 8. 结果保存函数 --------------------------
def save_results_to_excel(solution, solution_fitness, filename="result1.xlsx"):
    drone_speed, theta, t_drop1, t_fuse1, t_drop2, t_fuse2, t_drop3, t_fuse3 = solution

    # 计算投放点和起爆点
    drop_points = []
    explosion_points = []

    for i, (t_drop, t_fuse) in enumerate([(t_drop1, t_fuse1), (t_drop2, t_fuse2), (t_drop3, t_fuse3)]):
        drop_x = drone_FY1[0] + drone_speed * math.cos(theta) * t_drop
        drop_y = drone_FY1[1] + drone_speed * math.sin(theta) * t_drop
        drop_z = drone_FY1[2]
        drop_points.append((drop_x, drop_y, drop_z))

        explosion_point = get_smoke_explosion_center(drone_FY1, drone_speed, theta, t_drop, t_fuse)
        explosion_points.append((explosion_point[0], explosion_point[1], explosion_point[2]))

    # 方向角度（0~360度，逆时针为正，x轴为正向）
    theta_deg = (theta * 180 / math.pi) % 360

    # 创建DataFrame
    data = {
        "无人机编号": ["FY1"] * 3,
        "飞行速度(m/s)": [drone_speed] * 3,
        "飞行方向(°)": [theta_deg] * 3,
        "弹药序号": [1, 2, 3],
        "投放时间(s)": [t_drop1, t_drop2, t_drop3],
        "投放点坐标X(m)": [p[0] for p in drop_points],
        "投放点坐标Y(m)": [p[1] for p in drop_points],
        "投放点坐标Z(m)": [p[2] for p in drop_points],
        "起爆时间(s)": [t_fuse1, t_fuse2, t_fuse3],
        "起爆点坐标X(m)": [p[0] for p in explosion_points],
        "起爆点坐标Y(m)": [p[1] for p in explosion_points],
        "起爆点坐标Z(m)": [p[2] for p in explosion_points],
        "总遮蔽时长(s)": [solution_fitness] * 3
    }

    df = pd.DataFrame(data)
    df.to_excel(filename, index=False)
    print(f"结果已保存到 {filename}")


# -------------------------- 9. 主程序 --------------------------
if __name__ == "__main__":
    # 设置遗传算法参数
    gene_space = [
        {'low': 70, 'high': 140},  # 速度
        {'low': 0, 'high': 2 * math.pi},  # 方向角
        {'low': 0, 'high': missile_time_to_fake},  # 投放延迟1
        {'low': 0, 'high': missile_time_to_fake},  # 起爆延迟1
        {'low': 0, 'high': missile_time_to_fake},  # 投放延迟2
        {'low': 0, 'high': missile_time_to_fake},  # 起爆延迟2
        {'low': 0, 'high': missile_time_to_fake},  # 投放延迟3
        {'low': 0, 'high': missile_time_to_fake}  # 起爆延迟3
    ]

    # 创建初始种群
    initial_population = create_starting_population()

    ga_instance = pygad.GA(
        num_generations=200,  # 增加迭代次数
        num_parents_mating=30,  # 增加父代数量
        fitness_func=fitness_func,
        sol_per_pop=100,  # 增加种群大小
        num_genes=8,
        gene_space=gene_space,
        initial_population=initial_population,
        parent_selection_type="tournament",
        crossover_type="two_points",
        mutation_type="random",
        mutation_percent_genes=15,  # 增加变异概率
        mutation_num_genes=2,
        keep_parents=10,  # 保留更多父代
        parallel_processing=4,
        on_generation=on_generation
    )

    print("开始遗传算法优化...")
    ga_instance.run()

    # 结果解析
    solution, solution_fitness, solution_idx = ga_instance.best_solution()
    drone_speed_opt, theta_opt, t_drop1_opt, t_fuse1_opt, t_drop2_opt, t_fuse2_opt, t_drop3_opt, t_fuse3_opt = solution

    # 方向角度（0~360度，逆时针为正，x轴为正向）
    theta_deg = (theta_opt * 180 / math.pi) % 360

    print("\n最优解：")
    print(f"1. 飞行参数：速度={drone_speed_opt:.2f}m/s, 方向角度={theta_deg:.2f}°")
    print(f"2. 第一枚弹：投放延迟={t_drop1_opt:.2f}s, 起爆延迟={t_fuse1_opt:.2f}s")
    print(f"3. 第二枚弹：投放延迟={t_drop2_opt:.2f}s, 起爆延迟={t_fuse2_opt:.2f}s")
    print(f"4. 第三枚弹：投放延迟={t_drop3_opt:.2f}s, 起爆延迟={t_fuse3_opt:.2f}s")
    print(f"5. 总遮蔽时长：{solution_fitness:.2f}s")

    # 保存结果到Excel
    save_results_to_excel(solution, solution_fitness)

    # 绘制迭代过程图
    plt.figure()
    plt.plot(range(1, len(generation_best_fitness) + 1), generation_best_fitness, marker='o')
    plt.xlabel("迭代次数")
    plt.ylabel("最大适应度")
    plt.title("遗传算法迭代过程")
    plt.grid(True)
    plt.savefig("optimization_process.png")
    plt.show()

    # 可视化三枚弹的遮蔽时间区间
    plt.figure(figsize=(10, 6))

    # 计算每枚弹的有效时间区间
    detonate_times = [
        t_drop1_opt + t_fuse1_opt,
        t_drop2_opt + t_fuse2_opt,
        t_drop3_opt + t_fuse3_opt
    ]

    colors = ['red', 'green', 'blue']
    labels = ['弹1', '弹2', '弹3']

    for i, detonate_t in enumerate(detonate_times):
        start_t = detonate_t
        end_t = detonate_t + smoke_valid_time

        # 采样判断每个时间点是否遮蔽
        t_values = np.arange(start_t, end_t, 0.1)
        occlusion_values = []

        for t in t_values:
            if is_effective_occlusion(t, drone_speed_opt, theta_opt,
                                      [t_drop1_opt, t_drop2_opt, t_drop3_opt][i],
                                      [t_fuse1_opt, t_fuse2_opt, t_fuse3_opt][i]):
                occlusion_values.append(1)
            else:
                occlusion_values.append(0)

        plt.plot(t_values, np.array(occlusion_values) + i * 0.1, color=colors[i], label=labels[i])

    plt.xlabel("时间 (s)")
    plt.ylabel("遮蔽状态")
    plt.yticks([0.05, 1.05, 2.05], labels)
    plt.title("三枚烟幕干扰弹的遮蔽时间区间")
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.savefig("smoke_occlusion_intervals.png")
    plt.show()